import net from 'net';
import dgram from 'dgram';
import {ProxyEntity} from "@/entityType/ProxyEntity.ts";

/**
 * 创建TCP代理
 * @param proxy 代理配置对象
 */
export function createTcpProxy(proxy: ProxyEntity) {
    // 创建TCP服务器
    const server = net.createServer((clientSocket) => {
        console.log(`[${new Date().toISOString()}] TCP代理 | 客户端 ${clientSocket.remoteAddress}:${clientSocket.remotePort} 连接至本地端口 ${proxy.localPort}`);

        // 连接目标服务器
        const targetSocket = net.connect({
            host: proxy.target,
            port: proxy.targetPort
        }, () => {
            console.log(`[${new Date().toISOString()}] TCP代理 | 已建立到目标 ${proxy.target}:${proxy.targetPort} 的连接`);
        });

        // 双向数据转发（客户端 -> 目标）
        clientSocket.pipe(targetSocket)
            .on('error', (err) => {
                console.error(`[${new Date().toISOString()}] TCP代理 | 目标转发错误: ${err.message}`);
            });

        // 双向数据转发（目标 -> 客户端）
        targetSocket.pipe(clientSocket)
            .on('error', (err) => {
                console.error(`[${new Date().toISOString()}] TCP代理 | 客户端转发错误: ${err.message}`);
            });

        // 客户端断开处理
        clientSocket.on('end', () => {
            console.log(`[${new Date().toISOString()}] TCP代理 | 客户端 ${clientSocket.remoteAddress}:${clientSocket.remotePort} 断开连接`);
            targetSocket.end();
        });

        // 目标断开处理
        targetSocket.on('end', () => {
            console.log(`[${new Date().toISOString()}] TCP代理 | 目标 ${proxy.target}:${proxy.targetPort} 连接关闭`);
            clientSocket.end();
        });

        // 错误处理
        clientSocket.on('error', (err) => {
            console.error(`[${new Date().toISOString()}] TCP代理 | 客户端错误: ${err.message}`);
            targetSocket.destroy();
        });

        targetSocket.on('error', (err) => {
            console.error(`[${new Date().toISOString()}] TCP代理 | 目标错误: ${err.message}`);
            clientSocket.destroy();
        });
    });

    return new Promise((resolve, reject) => {
        // 启动TCP服务器
        server.listen({
            port: proxy.localPort,
            host: '0.0.0.0'  // 监听所有网络接口
        }, () => {
            console.log(`[${new Date().toISOString()}] TCP代理 | 成功启动，本地端口 ${proxy.localPort} -> 目标 ${proxy.target}:${proxy.targetPort}`);
            resolve(server);
        });

        // 服务器错误处理
        server.on('error', (err) => {
            console.error(`[${new Date().toISOString()}] TCP代理 | 启动失败（端口 ${proxy.localPort}）: ${err.message}`);
            // @ts-ignore
            if (err.code === 'EADDRINUSE') {
                console.error(`[${new Date().toISOString()}] TCP代理 | 端口 ${proxy.localPort} 已被占用`);
                reject(`[${new Date().toISOString()}] TCP代理 | 端口 ${proxy.localPort} 已被占用`)
            }
        });

        server.on('close', () => {
            console.error(`[${new Date().toISOString()}] TCP代理 | 停止代理完成（端口 ${proxy.localPort}）`);
        });
    })
}

/**
 * 创建UDP代理（基础实现，适用于单客户端或简单场景）
 * @param proxy 代理配置对象
 */
export function createUdpProxy(proxy: ProxyEntity) {
    return new Promise((resolve, reject) => {
        const socket = dgram.createSocket({type: 'udp4', reuseAddr: true});
        const requestMap = new Map();  // 存储请求标识与客户端地址的映射（示例使用简单计数）

        // 生成唯一请求ID（示例用简单递增ID，实际建议用更可靠的方案）
        let requestId = 0;

        // 接收消息处理
        socket.on('message', (msg, rinfo) => {
            const targetKey = `${proxy.target}:${proxy.targetPort}`;
            const isRequest = !requestMap.has(targetKey);

            if (isRequest) {
                // 处理客户端请求（假设首次接收的是客户端请求）
                console.log(`[${new Date().toISOString()}] UDP代理 | 收到客户端 ${rinfo.address}:${rinfo.port} 请求，转发至 ${targetKey}`);

                // 生成请求ID并存储客户端信息
                const currentRequestId = requestId++;
                requestMap.set(targetKey, {
                    clientId: rinfo,
                    requestId: currentRequestId
                });

                // 转发请求到目标（添加简单请求头标识）
                const requestWithId = Buffer.alloc(4 + msg.length);
                requestWithId.writeUInt32BE(currentRequestId, 0);  // 前4字节为请求ID
                // @ts-ignore
                msg.copy(requestWithId, 4);
                // @ts-ignore
                socket.send(requestWithId, 0, requestWithId.length, proxy.targetPort, proxy.target, (err) => {
                    if (err) {
                        console.error(`[${new Date().toISOString()}] UDP代理 | 转发请求失败: ${err.message}`);
                        requestMap.delete(targetKey);
                    }
                });
            } else {
                // 处理目标响应（假设后续接收的是目标响应）
                console.log(`[${new Date().toISOString()}] UDP代理 | 收到目标 ${targetKey} 响应，尝试回传客户端`);

                // 解析请求ID（前4字节）
                const receivedRequestId = msg.readUInt32BE(0);
                const stored = requestMap.get(targetKey);

                if (stored && stored.requestId === receivedRequestId) {
                    // 匹配到对应客户端，回传数据（去除请求头）
                    const responseData = msg.slice(4);
                    // @ts-ignore
                    socket.send(responseData, 0, responseData.length, stored.clientId.port, stored.clientId.address, (err) => {
                        if (err) {
                            console.error(`[${new Date().toISOString()}] UDP代理 | 回传响应失败: ${err.message}`);
                        }
                        requestMap.delete(targetKey);  // 清理已完成请求
                    });
                } else {
                    console.warn(`[${new Date().toISOString()}] UDP代理 | 未找到匹配的请求ID（${receivedRequestId}），可能为无效响应`);
                }
            }
        });

        // 错误处理
        socket.on('error', (err) => {
            console.error(`[${new Date().toISOString()}] UDP代理 | 错误: ${err.message}`);
            socket.close();
            reject(`[${new Date().toISOString()}] UDP代理 | 错误: ${err.message}`)
        });

        socket.on('close', () => {
            console.error(`[${new Date().toISOString()}] UDP代理 | 停止代理完成（端口 ${proxy.localPort}）`);
        });

        // 启动UDP服务
        socket.bind({
            port: proxy.localPort,
            address: '0.0.0.0'
        }, () => {
            console.log(`[${new Date().toISOString()}] UDP代理 | 成功启动，本地端口 ${proxy.localPort} -> 目标 ${proxy.target}:${proxy.targetPort}`);
            resolve(socket)
        });
    })
}